from DNA import DNA_lcq
from GeneticAlgorithm import GeneticAlgorithm
import matplotlib.pyplot as plt  
import numpy as np

def fx1(x, array):
    a, b, c = array
    return a * x ** 2 + b * x + c

# 下叉下降 右转
# weight = np.array([0, 30, 60, 90, 120, 150, 180, 210, 240]) 
# values = np.array([13, 19.5, 24.5, 31.5, 38.5, 45,54, 62, 70])

# 下叉下降 左转
weight = np.array([0, 30, 60, 90, 120, 150, 180, 210, 240, 270]) 
values = np.array([-1.5, -1.5, 3, 10.5, 18.5, 23.5, 31, 39, 46.5, 57])

# 安装架下降 右转
# weight = np.array([90, 120, 150, 180, 210, 240]) 
# values = np.array([5, 6, 7, 8, 10, 10])

# 安装架下降 左转
# weight = np.array([90, 120, 150, 180, 210, 240]) 
# values = np.array([2, 2, 3.5, 4, 5, 6.5])

# 适应度函数

def fun(array):
    # 接收一个 1 * n 的列表
    # array : [a, b, c]
    # a, b, c
    # a * x ** 2 + b * x + c
    # 返回一个数字，最大值优化

    # 要都大于数据集
    values_fit = fx1(weight, array)  # 预测值

    error_total = 0  # 误差
    for value, value_fit in zip(values, values_fit):
        error = abs(value_fit - value)
        error_total += error * error  # 平方差求和
        if value_fit < value + 1:     # 如果补偿值不大于下降值1mm，惩罚他3000分(3000分随便给的)
            error_total += 3000

    a, b, c = array
    if b < 2 * a:    # 如果对称轴不在y轴左侧，惩罚500分（500分随便给的）
        error_total += 500

    return -error_total

# 超参数
# goal_fun:             目标函数 (默认求最大值)
# parameter_range:      参数的约束范围（左右闭区间）
# parameter_accuracy:   各个参数精度（总共多少种可能）
# crossover_p:          交叉概率
# variation_p:          变异概率
# coding_mode:          编码方式
# max_epoch:            最大迭代次数
# population_size:      种群个数


n = 3  # 尺寸
left = -100  # 下限（闭区间）
right = 100  # 上限 （闭区间）

parameter = {"goal_fun": fun,
             "parameter_range": [[0, right], [left, right], [left, right]],
             "parameter_accuracy": None,
             "crossover_p": 0.7,
             "variation_p": 0.5,
             "coding_mode": DNA_lcq,
             "max_epoch": 500,
             "population_size": 15}

# 这里默认是最大值优化，如果需要最小值优化，则需要对fun函数的结果乘以 -1
Gene = GeneticAlgorithm(parameter)
ans = Gene.go(log=True)  # log为False则不打印日志
print("最优解为：", ans)  # 走！
print("对应结果为：", fun(ans))



# 绘制结果  
x_fit = np.linspace(0, 300, 100)  
plt.scatter(weight, values, label='Data', color='black')  


plt.plot(x_fit, fx1(x_fit, ans), label='Linear Fit', color='red')  

plt.xlabel('Weight (kg)')  
plt.ylabel('Value')  
plt.title('Data Fitting')  
plt.legend()  
plt.grid()  
plt.show()